import abc
from typing import Any
from langchain.tools import BaseTool

import logging
import sys
import os
import pandas as pd
import utils.config_util as utils
from langchain_openai import ChatOpenAI
from langchain_community.embeddings.openai import OpenAIEmbeddings
import sqlite3

import pandas as pd
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer

from utils.instances import TOKENIZER, LLM

g_example_question_list = list()
g_example_sql_list = list()
g_example_fa_list = list()
g_example_info_list = list()
g_example_token_list = list()

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

class FinSQLRAG(BaseTool, abc.ABC):
    # 添加类型注解
    name: str = "查询金融数据库"
    description: str = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 用例中没有用到 arun 不予具体实现
        pass

    def _run(self, para) -> str:
        query = para
        try:
            if len(g_example_question_list) <= 0:
                sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
                for cyc in range(len(sql_examples_file)):
                    g_example_question_list.append(sql_examples_file[cyc:cyc + 1]['问题'][cyc])
                    g_example_sql_list.append(sql_examples_file[cyc:cyc + 1]['SQL'][cyc])
                    g_example_info_list.append(sql_examples_file[cyc:cyc + 1]['资料'][cyc])
                    g_example_fa_list.append(sql_examples_file[cyc:cyc + 1]['FA'][cyc])
                    tokens = TOKENIZER(sql_examples_file[cyc:cyc + 1]['问题'][cyc])
                    tokens = tokens['input_ids']
                    g_example_token_list.append(tokens)
            result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list,
                                              g_example_token_list)
            print("query sql is:", sql)
            conn = sqlite3.connect(
                r'C:\Users\张朋飞\Desktop\数字人\数字人\Fay-fay-agent-edition0830\agent\tools\data\博金杯比赛数据.db')
            cs = conn.cursor()
            success_flag, exc_result = query_db(sql, cs)
            conn.close()
            answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list,
                                     g_example_fa_list,
                                     g_example_token_list)
            return answer
        except Exception as e:
            print("Retriever Error:", e)
            return ('FinSQLRAG处理异常！')


if __name__ == "__main__":
    tool = FinSQLRAG()
    result = tool.run("请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。")
    print(result)